import os
import pandas as pd
import xmind
import shutil
import zipfile
import rarfile
from datetime import datetime
import argparse

priority_dict = {"高": 'priority-1', "中": 'priority-2', "低": 'priority-3'}


def extract(d_path, f_path, mode="zip"):
    """
    zip解压缩乱码问题处理
    :param d_path:
    :param f_path:
    :return:
    """
    root = d_path
    if not os.path.exists(root):
        os.makedirs(root)

    if mode == 'zip':
        zf = zipfile.ZipFile(f_path, "r")
    elif mode == 'rar':
        zf = rarfile.RarFile(f_path, "r")

    for n in zf.infolist():
        srcName = n.filename
        try:
            decodeName = srcName.encode("cp437").decode("utf-8")
        except:
            try:
                decodeName = srcName.encode("cp437").decode("gbk")
            except:
                decodeName = srcName
        spiltArr = decodeName.split("\\")
        path = root
        for temp in spiltArr:
            path = os.path.join(path, temp)

        if decodeName.endswith("\\"):
            if not os.path.exists(path):
                os.makedirs(path)
        else:
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            f = open(path, "wb")
            f.write(zf.read(srcName))
            f.close()
    zf.close()


def aftertreatment(path):
    """
    **解决 xmind8 可以打开 xmind2020 报错
    """
    # 修改名字
    path = path.strip()
    retval = os.path.dirname(os.path.abspath(__file__))
    folder = os.path.dirname(path)
    if not folder:
        folder = retval
    name = os.path.basename(path)
    unzip_folder = os.path.splitext(name)[0]
    zip_name = unzip_folder + ".zip"
    os.chdir(folder)
    os.rename(name, zip_name)
    os.chdir(retval)
    # 解压
    unzip_path = os.path.join(folder, unzip_folder)
    if not os.path.exists(unzip_path):
        os.mkdir(unzip_path)

    inf_folder = os.path.join(unzip_path, "META-INF")
    if not os.path.exists(inf_folder):
        os.mkdir(inf_folder)

    extract(unzip_path, os.path.join(folder, zip_name))
    # 2020以后格式中有manifest.xml，需要以下代码建立这个文件
    xmlfilepath = os.path.join(inf_folder, "manifest.xml")
    xml = '''<?xml version="1.0" encoding="UTF-8" standalone="no"?><manifest xmlns="urn:xmind:xmap:xmlns:manifest:1.0"><file-entry full-path="comments.xml" media-type=""/><file-entry full-path="content.xml" media-type="text/xml"/><file-entry full-path="markers/" media-type=""/><file-entry full-path="markers/markerSheet.xml" media-type=""/><file-entry full-path="META-INF/" media-type=""/><file-entry full-path="META-INF/manifest.xml" media-type="text/xml"/><file-entry full-path="meta.xml" media-type="text/xml"/><file-entry full-path="styles.xml" media-type=""/></manifest>'''
    with open(xmlfilepath, 'w', encoding='utf-8') as f:
        f.write(xml)

    os.remove(os.path.join(folder, zip_name))
    shutil.make_archive(unzip_path, 'zip', unzip_path)
    file_path = unzip_path + '.zip'
    os.chdir(os.path.dirname(file_path))
    os.rename(os.path.basename(file_path), name)
    os.chdir(retval)
    shutil.rmtree(unzip_path)


def create_case_node(case_df, case_parent_topic, case_nums):
    for _, row in case_df.iterrows():
        case_topic = case_parent_topic.addSubTopic()
        case_title = row['用例标题']
        case_topic.setTitle(f'[case] {case_title}')
        step_topic = case_topic.addSubTopic()
        precondition = row.get('前置条件', '')
        steps = row.get('测试步骤', '')
        expected_result = row.get('预期结果', '')
        step_topic.setTitle(f'[pre]\n{precondition}\n[step]\n{steps}')

        result_topic = step_topic.addSubTopic()
        result_topic.setTitle(f'[result]\n{expected_result}')

        # 添加标记
        is_smoke_case = row.get('是否冒烟用例', '否')
        priority = row.get('优先级', '').strip()
        if is_smoke_case == '是':
            case_topic.addMarker('flag-red')
        if priority:
            case_topic.addMarker(priority_dict.get(priority))
        case_nums += 1
    return case_nums


def case_group_by_type(module_df, module_topic, case_nums):
    case_type_groups = module_df.groupby('用例类型')
    for case_type, case_type_df in case_type_groups:
        if case_type:
            case_type_topic = module_topic.addSubTopic()
            case_type_topic.setTitle(f'[type] {case_type}')
            case_nums = create_case_node(case_type_df, case_type_topic, case_nums)
        else:
            case_nums = create_case_node(case_type_df, module_topic, case_nums)
    return case_nums


def create_xmind_from_excel(excel_file, xmind_file):
    try:
        # 初始化用例数量
        case_nums = 0

        # 读取 Excel 文件
        original_df = pd.read_excel(excel_file)
        df = original_df.fillna('')

        # 创建一个新的 XMind 工作簿
        template_path = 'template.xmind'
        workbook = xmind.load(template_path)

        sheet = workbook.getPrimarySheet()
        sheet.setTitle('用例信息')
        root_topic = sheet.getRootTopic()

        # 以项目版本作为根主题
        project_versions = df['版本名称'].unique()
        project = df['项目名称'].unique()
        root_topic.setTitle(project[0])
        for project_version in project_versions:
            project_topic = root_topic.addSubTopic()
            project_topic.setTitle(f"[version]{project_version}")

            # 筛选当前项目版本的数据
            project_df = df[df['版本名称'] == project_version]
            requirement_names = project_df['需求名称'].unique()
            for requirement_name in requirement_names:
                requirement_topic = project_topic.addSubTopic()
                requirement_topic.setTitle(f'[req.] {requirement_name}')

                # 筛选当前需求的数据
                requirement_df = project_df[project_df['需求名称'] == requirement_name]
                module_names = requirement_df[requirement_df['功能模块'] != '']['功能模块'].unique()
                for module_name in module_names:
                    module_topic = requirement_topic.addSubTopic()
                    module_topic.setTitle(f'[module] {module_name}')

                    # 筛选当前模块的数据
                    module_df = requirement_df[requirement_df['功能模块'] == module_name]
                    case_nums = case_group_by_type(module_df, module_topic, case_nums)

                # 处理功能模块为空的情况
                no_module_df = requirement_df[requirement_df['功能模块'] == '']
                if not no_module_df.empty:
                    case_nums = case_group_by_type(no_module_df, requirement_topic, case_nums)

        # 保存 XMind 文件
        if xmind_file:
            save_xmind = xmind_file.strip()
        else:
            # 获取当前脚本文件的绝对路径
            current_file_path = os.path.abspath(__file__)
            # 获取当前脚本文件所在的目录路径
            current_dir = os.path.join(os.path.dirname(current_file_path), 'xmind_file')
            current_time = datetime.now().strftime("%Y%m%d")
            excel_name = os.path.basename(excel_file).split('.')[0]
            save_xmind = os.path.join(current_dir, f'xmind_{excel_name}_{current_time}.xmind')

        if os.path.exists(save_xmind):
            try:
                # 尝试删除文件
                os.remove(save_xmind)
                print(f'{save_xmind} 文件已成功删除。')
            except OSError as e:
                # 若删除过程中出现错误，打印错误信息
                print(f'删除 {save_xmind} 文件时出现错误: {e}')
        xmind.save(workbook, path=save_xmind)
        aftertreatment(save_xmind)
        return f"【EXCEL转XMIND】成功将 {excel_file} 转换为 {save_xmind}, 共 {case_nums} 条用例"
    except Exception as e:
        print(e)
        return f"【EXCEL转XMIND】转换{excel_file}失败"


if __name__ == "__main__":
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='将 Excel 文件转换为 XMind 文件')

    # 添加命令行参数
    parser.add_argument('excel_file', type=str, help='输入的 excel 文件路径')
    parser.add_argument('-o', '--output', type=str,
                        help='输出的 xmind 文件路径，非必填，默认为 xmind{excel_file}_{%Y%m%d}.xmind')

    # 解析命令行参数
    args = parser.parse_args()

    # 获取传入的参数值
    excel_file = args.excel_file
    xmind_file = args.output

    # 调用转换函数
    result = create_xmind_from_excel(excel_file, xmind_file)
    print(result)